import numpy as np
from numpy.random import choice
import random

def my_side_slope_obs_terrain(terrain, 
                    platform_len=2.5, 
                    platform_height=0., 
                    num_stones=8, 
                    x_range=[1.8, 1.9],  # d
                    y_range=[0., 0.1],   # d
                    z_range=[-0.2, 0.2],  
                    stone_len=1.0,  # d
                    stone_width=0.6,
                    pad_width=0.1,
                    pad_height=0.5,
                    incline_height=0.1,  # d
                    last_incline_height=0.6,  # d
                    last_stone_len=1.6,
                    pit_depth=[0.5, 1.],
                    num_obstacles=3,
                    obstacle_interval=0,
                    obstacle_size=[0.1, 0.5],
                    obstacle_height=1.,
                    avoid_distance=0.8):
    goals = np.zeros((num_stones+2+6, 3+3))
    terrain.height_field_raw[:] = -round(np.random.uniform(pit_depth[0], pit_depth[1]) / terrain.vertical_scale)
    
    mid_y = terrain.length // 2  # length is actually y width
    stone_len = np.random.uniform(*stone_len)
    stone_len = 2 * round(stone_len / 2.0, 1)
    stone_len = round(stone_len / terrain.horizontal_scale)
    dis_x_min = stone_len + round(x_range[0] / terrain.horizontal_scale)
    dis_x_max = stone_len + round(x_range[1] / terrain.horizontal_scale)
    dis_y_min = round(y_range[0] / terrain.horizontal_scale)
    dis_y_max = round(y_range[1] / terrain.horizontal_scale)

    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    terrain.height_field_raw[0:platform_len, :] = platform_height

    stone_width = round(stone_width / terrain.horizontal_scale)
    last_stone_len = round(last_stone_len / terrain.horizontal_scale)

    incline_height = round(incline_height / terrain.vertical_scale)
    last_incline_height = round(last_incline_height / terrain.vertical_scale)

    dis_x = platform_len - np.random.randint(dis_x_min, dis_x_max) + stone_len // 2
    goals[0, :2] = [platform_len -  stone_len // 2, mid_y]
    left_right_flag = np.random.randint(0, 2)
    dis_z = 0
    
    for i in range(num_stones):
        dis_x += np.random.randint(dis_x_min, dis_x_max)
        pos_neg = round(2*(left_right_flag - 0.5))
        dis_y = mid_y + pos_neg * np.random.randint(dis_y_min, dis_y_max)
        if i == num_stones - 1:
            dis_x += last_stone_len // 4
            heights = np.tile(np.linspace(-last_incline_height, last_incline_height, stone_width), (last_stone_len, 1)) * pos_neg
            terrain.height_field_raw[dis_x-last_stone_len//2:dis_x+last_stone_len//2, dis_y-stone_width//2: dis_y+stone_width//2] = heights.astype(int) + dis_z
        else:
            heights = np.tile(np.linspace(-incline_height, incline_height, stone_width), (stone_len, 1)) * pos_neg
            terrain.height_field_raw[dis_x-stone_len//2:dis_x+stone_len//2, dis_y-stone_width//2: dis_y+stone_width//2] = heights.astype(int) + dis_z
        
        goals[i+1, :2] = [dis_x, dis_y]
        goals[i+1, 2] = 2
        left_right_flag = 1 - left_right_flag
        
    final_dis_x = dis_x + 2*np.random.randint(dis_x_min, dis_x_max)
    final_platform_start = dis_x + last_stone_len // 2 + round(0.05 // terrain.horizontal_scale)
    terrain.height_field_raw[final_platform_start:, :] = platform_height
    # goals[-1] = [final_dis_x, mid_y]
    goals[num_stones+1, :2] = [final_dis_x, mid_y]

    # creating obstacle
    obstacle_min = int(obstacle_size[0] / terrain.horizontal_scale)
    obstacle_max = int(obstacle_size[1] / terrain.horizontal_scale)
    # obstacle_max_range = terrain.height_field_raw.shape[0] - goals[num_stones+1][0]

    final_goal_x = terrain.height_field_raw.shape[0] - 25

    obstacle_x_centers = np.linspace(goals[num_stones+1][0], final_goal_x, 7)
    if num_obstacles == 1:
        obstacle_x_centers = obstacle_x_centers[3].reshape(1,)
    else:
        obstacle_x_centers = obstacle_x_centers[range(1,7,4)]
    obstacle_y_centers = mid_y
    obstacle_height = obstacle_height / terrain.vertical_scale
    avoid_distance = avoid_distance / terrain.horizontal_scale
    obstacle_goals = np.zeros((num_obstacles, 2+3))
    for i in range(num_obstacles):
        obstacle_width = np.random.randint(obstacle_min, obstacle_max)
        obstacle_length = np.random.randint(obstacle_min*1.5, obstacle_max*1.5) 
        obs_left_right = np.random.randint(0, 2)
        left_or_right = round(2*(obs_left_right - 0.5))
        y_offset = int(np.random.uniform(0, 1) * obstacle_min // 2) * left_or_right
        y_center = y_offset + obstacle_y_centers
        x_center = int(obstacle_x_centers[i])
        obstacle_goals[i, :2] = [x_center, y_center + ((-left_or_right) * (avoid_distance+(obstacle_length // 2)))]
        # obstacle_goals[i, 2] = 6
        terrain.height_field_raw[x_center-obstacle_width//2:x_center+obstacle_width//2, 
                                 y_center-obstacle_length//2:y_center+obstacle_length//2] = obstacle_height
        
    # creating goals
    last_goals = np.vstack((np.linspace(goals[num_stones+1][0], final_goal_x, 6, endpoint=False), np.array([mid_y]).repeat(6))).transpose(1,0)[1:,:]
    if num_obstacles == 1:
        last_goals[2] = obstacle_goals[0, :2]
    elif num_obstacles ==2:
        last_goals[0] = obstacle_goals[0, :2]
        last_goals[4] = obstacle_goals[1, :2]
    goals[num_stones+2:-1,:2] = last_goals[:,:]
    goals[-1:,:2] = [final_goal_x, mid_y]
    terrain.goals = goals * terrain.horizontal_scale
    
    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height
    # for item in goals:
    #     terrain.height_field_raw[int(item[0]-2):int(item[0]+2), int(item[1]-2):int(item[1]+2)] = 500

def my_hurdle_obs_terrain(terrain,
                            platform_len=2.5, 
                            platform_height=0., 
                            num_stones=8,
                            stone_len=0.3,
                            x_range=[1.5, 2.4],
                            y_range=[-0.4, 0.4],
                            hurdle_width_range=[0.4, 0.8],
                            hurdle_height_range=[0.2, 0.3],
                            pad_width=0.1,
                            pad_height=0.5,
                            num_obstacles=3,
                            obstacle_interval=0,
                            obstacle_size=[0.1, 0.5],
                            obstacle_height=1.,
                            avoid_distance=0.8,
                            flat=False):
    goals = np.zeros((num_stones+2+6, 3+4))
    
    mid_y = terrain.length // 2  # length is actually y width

    dis_x_min = round(x_range[0] / terrain.horizontal_scale)
    dis_x_max = round(x_range[1] / terrain.horizontal_scale)
    dis_y_min = round(y_range[0] / terrain.horizontal_scale)
    dis_y_max = round(y_range[1] / terrain.horizontal_scale)

    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    terrain.height_field_raw[0:platform_len, :] = platform_height

    stone_len = round(stone_len / terrain.horizontal_scale)

    dis_x = platform_len
    goals[0, :2] = [platform_len - 1, mid_y]

    if num_obstacles == 3:
        interval_index = np.linspace(1, 5, num_obstacles+2)[range(0,5,2)]
    else:
        interval_index = np.linspace(1, 5, num_obstacles*2, endpoint=False)[range(1, num_obstacles*2, 2)]
    obstacle_min = int(obstacle_size[0] / terrain.horizontal_scale)
    obstacle_max = int(obstacle_size[1] / terrain.horizontal_scale)
    obstacle_height = obstacle_height / terrain.vertical_scale
    avoid_distance = avoid_distance / terrain.horizontal_scale
    obstacle_interval = 3.5 / terrain.horizontal_scale
    goals_index = 1
    for i in range(num_stones):
        rand_x = np.random.randint(dis_x_min, dis_x_max)
        rand_y = np.random.randint(dis_y_min, dis_y_max)
        if i in interval_index:
            dis_x += int(obstacle_interval)
            # add obstacles and add goals 
        dis_x += rand_x
        if not flat:
            hurdle_height = np.random.uniform(*hurdle_height_range) if isinstance(hurdle_height_range, (tuple, list)) else hurdle_height_range
            hurdle_height /= terrain.vertical_scale
            hurdle_width = np.random.uniform(*hurdle_width_range) if isinstance(hurdle_width_range, (tuple, list)) else hurdle_width_range
            hurdle_width /= terrain.horizontal_scale; hurdle_width = round(hurdle_width//2)
            terrain.height_field_raw[dis_x-stone_len//2:dis_x+stone_len//2, ] = hurdle_height
            terrain.height_field_raw[dis_x-stone_len//2:dis_x+stone_len//2, :mid_y+rand_y-hurdle_width] = 0
            terrain.height_field_raw[dis_x-stone_len//2:dis_x+stone_len//2, mid_y+rand_y+hurdle_width:] = 0
        if i in interval_index:
            goals[goals_index, :2] = [last_dis_x+rand_x//2, goals[goals_index-1][1]]
            goals[goals_index, 3] = 2
            goals[goals_index, 4] = last_hurdle_height * terrain.vertical_scale
            goals[goals_index, 5] = 2 * last_hurdle_width * terrain.horizontal_scale
            goals_index += 1
            x_center = int((dis_x-rand_x//2 + last_dis_x+rand_x//2) // 2)
            obstacle_y_centers = int((mid_y + rand_y + goals[goals_index-1][1]) // 2)
            obs_left_right = np.random.randint(0, 2)
            left_or_right = round(2*(obs_left_right - 0.5))
            y_offset = int(np.random.uniform(0, 1) * obstacle_min // 2) * left_or_right
            y_center = y_offset + obstacle_y_centers
            obstacle_width = np.random.randint(obstacle_min, obstacle_max)
            obstacle_length = np.random.randint(obstacle_min, obstacle_max) 
            # goals[goals_index] = [x_center, y_center + (obstacle_length // 2) * (-left_or_right) + ((-left_or_right) * avoid_distance)]
            goals[goals_index, :2] = [x_center, y_center  + ((-left_or_right) * (avoid_distance+(obstacle_length // 2)))]
            # goals[goals_index, 3] = 6
            goals_index += 1
            terrain.height_field_raw[x_center-obstacle_width//2:x_center+obstacle_width//2, 
                                 y_center-obstacle_length//2:y_center+obstacle_length//2] = obstacle_height
            
        last_dis_x = dis_x
        last_mid_y = mid_y + rand_y
        goals[goals_index, :2] = [dis_x-rand_x//2, mid_y + rand_y]
        if i != 0:
            goals[goals_index, 3] = 2
            goals[goals_index, 4] = last_hurdle_height * terrain.vertical_scale
            goals[goals_index, 5] = 2 * last_hurdle_width * terrain.horizontal_scale
        # if i in interval_index:
        #     goals[goals_index, 3] = 6
        goals_index += 1
        last_hurdle_height = hurdle_height
        last_hurdle_width = hurdle_width
    
    final_dis_x = dis_x + 20
    
    goals[goals_index, :2] = [final_dis_x, mid_y]
    goals[goals_index, 3] = 2
    goals[goals_index, 4] = hurdle_height * terrain.vertical_scale
    goals[goals_index, 5] = hurdle_width * terrain.horizontal_scale
    goals_index += 1
    # create rest goals
    if num_obstacles != 3:
        last_obstacle_x = int(np.linspace(final_dis_x, terrain.height_field_raw.shape[0], (3-num_obstacles)*2+1, endpoint=False)[(3-num_obstacles)])
        obstacle_width = np.random.randint(obstacle_min, obstacle_max)
        obstacle_length = np.random.randint(obstacle_min, obstacle_max) 
        obs_left_right = np.random.randint(0, 2)
        left_or_right = round(2*(obs_left_right - 0.5))
        y_offset = int(np.random.uniform(0, 1) * obstacle_min // 2) * left_or_right
        last_obstacle_y = y_offset + mid_y
        
        terrain.height_field_raw[last_obstacle_x-obstacle_width//2:last_obstacle_x+obstacle_width//2, 
                                 last_obstacle_y-obstacle_length//2:last_obstacle_y+obstacle_length//2] = obstacle_height
        last_goals = np.vstack((np.linspace(final_dis_x, terrain.height_field_raw.shape[0], (3-num_obstacles)*2+1, endpoint=False), 
                                np.array([mid_y]).repeat((3-num_obstacles)*2+1))).transpose(1,0)[1:,:]
        last_goals[(3-num_obstacles)-1, :] = \
            [last_obstacle_x, last_obstacle_y  + ((-left_or_right) * (avoid_distance + (obstacle_length // 2)))]
        last_goals = np.concatenate((last_goals, np.zeros(((3-num_obstacles)*2, 4))), axis=1)
        # last_goals[last_goals[:, 1]!=mid_y, 3] = 6
        # last_goals[np.roll((last_goals[:, 1]!=mid_y), 1), 3] = 6
        goals[goals_index:,:2] = last_goals[:, :2]
        goals[goals_index:,3] = last_goals[:,3]
    
    terrain.goals = np.concatenate((goals[:,:3] * terrain.horizontal_scale, goals[:, -4:]), axis=1)
    
    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height

def my_platform_obs_terrain(terrain,
                            platform_len=2.5, 
                            platform_height=0., 
                            num_stones=8,
                            x_range=[0.2, 0.4],
                            y_range=[-0.15, 0.15],
                            step_width_range=[0.45, 0.5],
                            step_height_range= 0.2,
                            pad_width=0.1,
                            pad_height=0.5,
                            num_prop=3,
                            num_obstacles=3,
                            obstacle_size=[0.1, 0.5],
                            obstacle_height=1.,
                            avoid_distance=0.5):
    goals = np.zeros((num_stones+2+6, 3+4))
    mid_y = terrain.length // 2  # length is actually y width

    dis_x_min = round((x_range[0] + step_height_range[1]) / terrain.horizontal_scale)
    dis_x_max = round((x_range[1] + step_height_range[1]) / terrain.horizontal_scale)
    dis_y_min = round(y_range[0] / terrain.horizontal_scale)
    dis_y_max = round(y_range[1] / terrain.horizontal_scale)
    
    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    terrain.height_field_raw[0:platform_len, :] = platform_height

    dis_x = platform_len
    last_dis_x = dis_x
    stair_height = 0
    goals[0, :2] = [platform_len - round(1 / terrain.horizontal_scale), mid_y]
    side_obstacle = num_obstacles >= 2
    mid_obstacle = (num_obstacles % 2 == 1)
    obstacle_min = int(obstacle_size[0] / terrain.horizontal_scale)
    obstacle_max = int(obstacle_size[1] / terrain.horizontal_scale)
    obstacle_height = obstacle_height / terrain.vertical_scale
    avoid_distance = avoid_distance / terrain.horizontal_scale
    goals_index = 1
    down_list = []
    for i in range(num_stones):
        step_height = np.random.uniform(*step_height_range) if isinstance(step_height_range, (tuple, list)) else step_height_range
        step_height /= terrain.vertical_scale
        step_width = np.random.uniform(*step_width_range) if isinstance(step_width_range, (tuple, list)) else step_width_range
        step_width /= terrain.horizontal_scale
        rand_x = np.random.randint(dis_x_min, dis_x_max)
        rand_y = np.random.randint(dis_y_min, dis_y_max)
        
        if i < num_stones // 2:
            stair_height += step_height
        elif i == num_stones // 2:
            down_list = np.random.uniform(0, stair_height, num_stones // 2 - 1)
            down_list = np.sort(down_list)[::-1]
            # print(down_list)
        elif i > num_stones // 2:
            step_height = stair_height - down_list[i-1-num_stones//2]
            stair_height = down_list[i-1-num_stones//2]
            # np.clip(stair_height, a_min=0, a_max=10000)
        
        if (i == 1 or i == 4) and side_obstacle:
            step_width *= 1.3
            rand_x = int(1.8*dis_x_max)
        
        if (i == 3 or i == 2) and mid_obstacle:
            step_width *= 1.3
            rand_x = int(1.3*dis_x_max)
            
        terrain.height_field_raw[dis_x:dis_x+rand_x, ] = stair_height
        dis_x += rand_x
        # terrain.height_field_raw[last_dis_x:dis_x, :np.clip(mid_y+rand_y-round(step_width//2), a_min=0, a_max=terrain.width)] = 0
        # terrain.height_field_raw[last_dis_x:dis_x, np.clip(mid_y+rand_y+round(step_width//2), a_min=0, a_max=terrain.width):] = 0
        terrain.height_field_raw[last_dis_x:dis_x, :mid_y+rand_y-round(step_width//2)] = 0
        terrain.height_field_raw[last_dis_x:dis_x, mid_y+rand_y+round(step_width//2):] = 0
        # terrain.height_field_raw[last_dis_x:dis_x, :] = 0
        # terrain.height_field_raw[last_dis_x:dis_x, mid_y+round(step_width//2):] = 0

        if (i == 1 or i == 4) and side_obstacle:
            obstacle_center_x = int(dis_x-rand_x//2)
            obstacle_center_y = int(mid_y+rand_y)
            obs_left_right = np.random.randint(0, 2)
            left_or_right = round(2*(obs_left_right - 0.5))
            y_offset = int(np.random.uniform(0, 1) * obstacle_min) * left_or_right
            y_center = y_offset + obstacle_center_y
            
            obstacle_width = np.random.randint(obstacle_min//2, obstacle_max//2)
            obstacle_length = np.random.randint(obstacle_min, obstacle_max) 
            obstacle_height_stair = obstacle_height + stair_height
            goals[goals_index, :2] = [(last_dis_x+obstacle_center_x)//2, obstacle_center_y]
            if i == 1:
                goals[goals_index, 3] = 3
                goals[goals_index, 4] = step_height * terrain.vertical_scale
                goals[goals_index, 5] = step_width * terrain.horizontal_scale   
            else:
                goals[goals_index, 3] = 4
                goals[goals_index, 4] = step_height * terrain.vertical_scale
                goals[goals_index, 5] = step_width * terrain.horizontal_scale 
            goals_index += 1
            goals[goals_index, :2] = [obstacle_center_x, y_center  + ((-left_or_right) * (avoid_distance+(obstacle_length // 2)))]
            # goals[goals_index, 3] = 6
            goals_index += 1
            goals[goals_index, :2] = [(dis_x+obstacle_center_x)//2, obstacle_center_y]
            # goals[goals_index, 3] = 6
            goals_index += 1
            terrain.height_field_raw[obstacle_center_x-obstacle_width//2:obstacle_center_x+obstacle_width//2, 
                                 y_center-obstacle_length//2:y_center+obstacle_length//2] = obstacle_height_stair
        elif i == 3 and mid_obstacle:
            obstacle_center_x = round((dis_x-rand_x//2 + goals[goals_index-1][0])/2)
            obstacle_center_y = round((goals[goals_index-1][1]+mid_y+rand_y) / 2)
            obs_left_right = np.random.randint(0, 2)
            left_or_right = round(2*(obs_left_right - 0.5))
            y_offset = int(np.random.uniform(0, 1) * obstacle_min) * left_or_right
            y_center = y_offset + obstacle_center_y
            
            obstacle_width = np.random.randint(obstacle_min//2, obstacle_max//2)
            obstacle_length = np.random.randint(obstacle_min, obstacle_max) 
            obstacle_height_stair = obstacle_height + stair_height
            goals[goals_index, :2] = [obstacle_center_x, y_center  + ((-left_or_right) * (avoid_distance+(obstacle_length // 2)))]
            # goals[goals_index, 3] = 6
            goals_index += 1
            goals[goals_index, :2] = [dis_x-rand_x//2, mid_y+rand_y]
            # goals[goals_index, 3] = 6
            goals_index += 1
            terrain.height_field_raw[obstacle_center_x-obstacle_width//2:obstacle_center_x+obstacle_width//2, 
                                 y_center-obstacle_length//2:y_center+obstacle_length//2] = obstacle_height_stair
        else:
            goals[goals_index, :2] = [dis_x-rand_x//2, mid_y+rand_y]
            if i<num_stones//2:
                goals[goals_index, 3] = 3
                goals[goals_index, 4] = step_height * terrain.vertical_scale
                goals[goals_index, 5] = step_width * terrain.horizontal_scale 
            else:
                if not(side_obstacle and i == num_stones//2):                                            
                    goals[goals_index, 3] = 4
                    goals[goals_index, 4] = step_height * terrain.vertical_scale
                    goals[goals_index, 5] = step_width * terrain.horizontal_scale 
            goals_index += 1
        
        last_dis_x = dis_x        
    
    final_dis_x = dis_x + 15
    if final_dis_x > terrain.width:
        final_dis_x = terrain.width - 0.5 // terrain.horizontal_scale
    goals[goals_index,:2] = [final_dis_x, mid_y]
    goals[goals_index, 3] = 4
    goals[goals_index, 4] = stair_height * terrain.vertical_scale
    goals[goals_index, 5] = step_width * terrain.horizontal_scale 
    
    goals_index += 1
    if num_obstacles == 3:
        goals[goals_index:,:2] = [terrain.height_field_raw.shape[0]-10, mid_y]
    else:
        last_obstacle_x = int(np.linspace(final_dis_x, terrain.height_field_raw.shape[0], int(14-goals_index)+1, endpoint=False)[int(((3-num_obstacles)*1.5))])
        obstacle_width = np.random.randint(obstacle_min, obstacle_max)
        obstacle_length = np.random.randint(obstacle_min, obstacle_max) 
        obs_left_right = np.random.randint(0, 2)
        left_or_right = round(2*(obs_left_right - 0.5))
        y_offset = int(np.random.uniform(0, 1) * obstacle_min // 2) * left_or_right
        last_obstacle_y = y_offset + mid_y
        terrain.height_field_raw[last_obstacle_x-obstacle_width//2:last_obstacle_x+obstacle_width//2, 
                                 last_obstacle_y-obstacle_length//2:last_obstacle_y+obstacle_length//2] = obstacle_height
        last_goals = np.vstack((np.linspace(final_dis_x, terrain.height_field_raw.shape[0], int(14-goals_index)+1, endpoint=False), 
                                np.array([mid_y]).repeat((int(14-goals_index)+1)))).transpose(1,0)[1:,:]
        last_goals[int(((3-num_obstacles)*1.5))-1, :] = [last_obstacle_x, last_obstacle_y  + ((-left_or_right) * (avoid_distance + (obstacle_length // 2)))]
        last_goals = np.concatenate((last_goals, np.zeros((int(14-goals_index), 4))), axis=1)
        # last_goals[last_goals[:, 1]!=mid_y, 3] = 6
        # last_goals[np.roll((last_goals[:, 1]!=mid_y), 1), 3] = 6
        goals[goals_index:,:2] = last_goals[:, :2]
        goals[goals_index:,3] = last_goals[:,3]
    # goals[goals_index:,:] = last_goals
    
    terrain.goals = np.concatenate((goals[:,:3] * terrain.horizontal_scale, goals[:, -4:]), axis=1)

    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height

def my_gap_obs_terrain(terrain,
                            platform_len=2.5, 
                            platform_height=0., 
                            num_gaps=6,
                            gap_size_range=0.3,
                            x_range=[1.6, 2.4],
                            y_range=[-1.2, 1.2],
                            stone_width_range=[0.6, 1.2],
                            gap_depth=-200,
                            pad_width=0.1,
                            pad_height=0.5,
                            num_obstacles=3,
                            obstacle_size=[0.1, 0.5],
                            obstacle_height=1.,
                            avoid_distance=0.8,
                            flat=False):
    goals = np.zeros((num_gaps+2+6, 3+4))
    
    mid_y = terrain.length // 2  # length is actually y width

    dis_y_min = round(y_range[0] / terrain.horizontal_scale)
    dis_y_max = round(y_range[1] / terrain.horizontal_scale)

    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    gap_depth = -np.random.uniform(*gap_depth) if isinstance(gap_depth, (tuple, list)) else -gap_depth
    gap_depth /= terrain.vertical_scale
        
    terrain.height_field_raw[0:platform_len, :] = platform_height

    dis_x = platform_len
    goals[0, :2] = [platform_len - 10, mid_y]
    last_dis_x = dis_x
    for i in range(num_gaps):
        stone_width = np.random.uniform(*stone_width_range) if isinstance(stone_width_range, (tuple, list)) else stone_width_range
        stone_width /= terrain.horizontal_scale
        
        gap_size = np.random.uniform(*gap_size_range) if isinstance(gap_size_range, (tuple, list)) else gap_size_range
        gap_size /= terrain.horizontal_scale
        rand_x = np.random.uniform(*x_range) if isinstance(x_range, (tuple, list)) else x_range
        rand_x /= terrain.horizontal_scale; rand_x += gap_size
        dis_x += rand_x ;dis_x = round(dis_x)
        rand_y = np.random.randint(dis_y_min, dis_y_max)
        if not flat:
            terrain.height_field_raw[round(dis_x-gap_size//2) : round(dis_x+gap_size//2), :] = gap_depth

        terrain.height_field_raw[last_dis_x:dis_x, :mid_y+rand_y-round(stone_width//2)] = gap_depth
        terrain.height_field_raw[last_dis_x:dis_x, mid_y+rand_y+round(stone_width//2):] = gap_depth
        
        last_dis_x = dis_x
        goals[i+1, :2] = [dis_x-rand_x//2, mid_y + rand_y]
        if i != 0:
            goals[i+1, 3] = 5
            goals[i+1, 6] = last_gap_size * terrain.horizontal_scale
        last_gap_size = gap_size
    final_dis_x = dis_x + 20
    
    if final_dis_x > terrain.width:
        final_dis_x = terrain.width - 0.5 // terrain.horizontal_scale

    goals[num_gaps+1, :2] = [final_dis_x, mid_y]
    goals[num_gaps+1, 3] = 5
    goals[i+1, 6] = last_gap_size
    # creating obstacle
    obstacle_min = int(obstacle_size[0] / terrain.horizontal_scale)
    obstacle_max = int(obstacle_size[1] / terrain.horizontal_scale)
    final_goal_x = terrain.height_field_raw.shape[0] - 25

    obstacle_x_centers = np.linspace(goals[num_gaps+1][0], final_goal_x, 7)
    if num_obstacles == 1:
        obstacle_x_centers = obstacle_x_centers[3].reshape(1,)
    else:
        obstacle_x_centers = obstacle_x_centers[range(1,7,4)]
    obstacle_y_centers = mid_y
    obstacle_height = obstacle_height / terrain.vertical_scale
    avoid_distance = avoid_distance / terrain.horizontal_scale
    obstacle_goals = np.zeros((num_obstacles, 2))
    for i in range(num_obstacles):
        obstacle_width = np.random.randint(obstacle_min, obstacle_max)
        obstacle_length = np.random.randint(obstacle_min*1.5, obstacle_max*1.5) 
        obs_left_right = np.random.randint(0, 2)
        left_or_right = round(2*(obs_left_right - 0.5))
        y_offset = int(np.random.uniform(0, 1) * obstacle_min // 2) * left_or_right
        y_center = y_offset + obstacle_y_centers
        x_center = int(obstacle_x_centers[i])
        obstacle_goals[i] = [x_center, y_center + ((-left_or_right) * (avoid_distance+(obstacle_length // 2)))]
        terrain.height_field_raw[x_center-obstacle_width//2:x_center+obstacle_width//2, 
                                 y_center-obstacle_length//2:y_center+obstacle_length//2] = obstacle_height
        
    # creating goals
    last_goals = np.vstack((np.linspace(goals[num_gaps+1][0], final_goal_x, 6, endpoint=False), np.array([mid_y]).repeat(6))).transpose(1,0)[1:,:]
    if num_obstacles == 1:
        last_goals[2] = obstacle_goals[0]
    elif num_obstacles ==2:
        last_goals[0] = obstacle_goals[0]
        last_goals[4] = obstacle_goals[1]
    
    last_goals = np.concatenate((last_goals, np.zeros((5, 3))), axis=1)
    # last_goals[last_goals[:, 1]!=mid_y, 3] = 6
    # last_goals[np.roll((last_goals[:, 1]!=mid_y), 1), 3] = 6
    goals[num_gaps+2:-1:,:2] = last_goals[:, :2]
    goals[num_gaps+2:-1:,3] = last_goals[:, 3]
    # goals[num_gaps+2:-1,:2] = last_goals[:,:]
    goals[-1:, :2] = [final_goal_x, mid_y]
    # goals[-1:, 3] = 6
    
    terrain.goals = np.concatenate((goals[:,:3] * terrain.horizontal_scale, goals[:, -4:]), axis=1)
    
    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height

def my_flat_terrain(terrain,
                    platform_len=2.5, 
                    platform_height=0., 
                    num_stones=8,
                    x_range=[1.5, 2.4],
                    y_range=[-0.4, 0.4],
                    pad_width=0.1,
                    pad_height=0.5,):
    goals = np.zeros((num_stones+2, 3+4))
    
    mid_y = terrain.length // 2  # length is actually y width

    dis_x_min = round(x_range[0] / terrain.horizontal_scale)
    dis_x_max = round(x_range[1] / terrain.horizontal_scale)
    dis_y_min = round(y_range[0] / terrain.horizontal_scale)
    dis_y_max = round(y_range[1] / terrain.horizontal_scale)
    
    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    terrain.height_field_raw[0:platform_len, :] = platform_height

    dis_x = platform_len
    goals[0, :2] = [platform_len - 1, mid_y]
    for i in range(num_stones):
        rand_x = np.random.randint(dis_x_min, dis_x_max)
        rand_y = np.random.randint(dis_y_min, dis_y_max)
        dis_x += rand_x
        goals[i+1, :2] = [dis_x-rand_x//2, mid_y + rand_y]
    final_dis_x = dis_x + np.random.randint(dis_x_min, dis_x_max)

    if final_dis_x > terrain.width:
        final_dis_x = terrain.width - 0.5 // terrain.horizontal_scale
    goals[-1, :2] = [final_dis_x, mid_y]
    
    terrain.goals = goals * terrain.horizontal_scale
    
    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height

def my_discrete_terrain(terrain,
                    platform_len=2.5, 
                    platform_height=0., 
                    num_stones=8,
                    x_range=[1.5, 2.4],
                    y_range=[-0.4, 0.4],
                    pad_width=0.1,
                    pad_height=0.5,
                    max_height=0.5,
                    min_size=0.2,
                    max_size=0.5,
                    num_rects=3,):
    goals = np.zeros((num_stones+2, 3+4))
    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    terrain.height_field_raw[0:platform_len, :] = platform_height
    
    dis_x_min = round(x_range[0] / terrain.horizontal_scale);dis_x_max = round(x_range[1] / terrain.horizontal_scale)
    dis_y_min = round(y_range[0] / terrain.horizontal_scale);dis_y_max = round(y_range[1] / terrain.horizontal_scale)
    mid_y = terrain.length // 2  # length is actually y width
    
    dis_x = platform_len
    goals[0, :2] = [platform_len - 1, mid_y]
    for i in range(num_stones):
        rand_x = np.random.randint(dis_x_min, dis_x_max)
        rand_y = np.random.randint(dis_y_min, dis_y_max)
        dis_x += rand_x
        goals[i+1, :2] = [dis_x-rand_x//2, mid_y + rand_y]
    final_dis_x = dis_x + np.random.randint(dis_x_min, dis_x_max)

    if final_dis_x > terrain.width:
        final_dis_x = terrain.width - 0.5 // terrain.horizontal_scale
    goals[-1, :2] = [final_dis_x, mid_y]
    
    terrain.goals = goals * terrain.horizontal_scale
    
    max_height = int(max_height / terrain.vertical_scale)
    min_size = int(min_size / terrain.horizontal_scale)
    max_size = int(max_size / terrain.horizontal_scale)

    (i, j) = terrain.height_field_raw.shape
    height_range = [-max_height, -max_height // 2, max_height // 2, max_height]
    width_range = range(min_size, max_size, 4)
    length_range = range(min_size, max_size, 4)

    for _ in range(num_rects):
        width = np.random.choice(width_range)
        length = np.random.choice(length_range)
        start_i = np.random.choice(range(0, i-width, 4))
        start_j = np.random.choice(range(0, j-length, 4))
        terrain.height_field_raw[start_i:start_i+width, start_j:start_j+length] = np.random.choice(height_range)
    
    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height

def my_slop_terrain(terrain,
                    platform_len=2.5, 
                    platform_height=0., 
                    num_stones=8,
                    x_range=[1.5, 2.4],
                    y_range=[-0.4, 0.4],
                    pad_width=0.1,
                    pad_height=0.5,
                    slop_end=[0.5, 0.9],
                    max_height=0.5,
                    down = False):
    slope = 1
    goals = np.zeros((num_stones+2, 3+4))
    
    mid_y = terrain.length // 2  # length is actually y width

    dis_x_min = round(x_range[0] / terrain.horizontal_scale)
    dis_x_max = round(x_range[1] / terrain.horizontal_scale)
    dis_y_min = round(y_range[0] / terrain.horizontal_scale)
    dis_y_max = round(y_range[1] / terrain.horizontal_scale)
    
    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    terrain.height_field_raw[0:platform_len, :] = platform_height

    dis_x = platform_len
    goals[0, :2] = [platform_len - 1, mid_y]
    for i in range(num_stones):
        rand_x = np.random.randint(dis_x_min, dis_x_max)
        rand_y = np.random.randint(dis_y_min, dis_y_max)
        dis_x += rand_x
        goals[i+1, :2] = [dis_x-rand_x//2, mid_y + rand_y]
    final_dis_x = dis_x + np.random.randint(dis_x_min, dis_x_max)

    if final_dis_x > terrain.width:
        final_dis_x = terrain.width - 0.5 // terrain.horizontal_scale
    goals[-1, :2] = [final_dis_x, mid_y]
    
    terrain.goals = goals * terrain.horizontal_scale
    slop_start = platform_len
    base_height = 0
    max_height = max_height / terrain.vertical_scale
    if down:
        max_height = -max_height
    for i, item in enumerate(slop_end):
        x = np.arange(0, int(((terrain.width)*item-slop_start)*2))
        y = np.arange(0, terrain.length)
        center_x = int((terrain.width)*item-slop_start)
        center_y = int(terrain.length / 2)
        xx, yy = np.meshgrid(x, y, sparse=True)
        xx = (center_x - np.abs(center_x-xx)) / center_x
        # yy = (center_y - np.abs(center_y-yy)) / center_y
        yy = (center_y - np.abs(center_y-yy)) / center_y
        xx = xx.reshape(int(((terrain.width)*item-slop_start)*2), 1)
        yy = yy.reshape(1, terrain.length)
        yy = 1.
        slope_array = (xx * yy)[:center_x, :]
        terrain.height_field_raw[int(slop_start):int((terrain.width)*item), :] += ((max_height * slope_array) + base_height).astype(terrain.height_field_raw.dtype)
        slop_start += (terrain.width)*item
        if i == len(slop_end) - 1: 
            terrain.height_field_raw[int((terrain.width)*item):, :] = max_height + base_height
        else:
            terrain.height_field_raw[int((terrain.width)*item):int(slop_start), :] = max_height + base_height
        base_height += max_height
    
    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height

def my_stairs_terrain(terrain,
                    step_width=0.3,
                    step_height=0.15,
                    platform_len=2.5, 
                    platform_height=0., 
                    num_stones=8,
                    x_range=[1.5, 2.4],
                    y_range=[-0.4, 0.4],
                    pad_width=0.1,
                    pad_height=0.5,
                    stairs_len=[0.3, 0.7],):
    goals = np.zeros((num_stones+2, 3+4))
    
    mid_y = terrain.length // 2  # length is actually y width

    dis_x_min = round(x_range[0] / terrain.horizontal_scale)
    dis_x_max = round(x_range[1] / terrain.horizontal_scale)
    dis_y_min = round(y_range[0] / terrain.horizontal_scale)
    dis_y_max = round(y_range[1] / terrain.horizontal_scale)
    
    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    terrain.height_field_raw[0:platform_len, :] = platform_height

    dis_x = platform_len
    goals[0, :2] = [platform_len - 1, mid_y]
    for i in range(num_stones):
        rand_x = np.random.randint(dis_x_min, dis_x_max)
        rand_y = np.random.randint(dis_y_min, dis_y_max)
        dis_x += rand_x
        goals[i+1, :2] = [dis_x-rand_x//2, mid_y + rand_y]
    final_dis_x = dis_x + np.random.randint(dis_x_min, dis_x_max)

    if final_dis_x > terrain.width:
        final_dis_x = terrain.width - 0.5 // terrain.horizontal_scale
    goals[-1, :2] = [final_dis_x, mid_y]
    
    terrain.goals = goals * terrain.horizontal_scale
    
    step_width = int(step_width / terrain.horizontal_scale)
    step_height = int(step_height / terrain.vertical_scale)

    height = 0
    start_x = round(platform_len * 1.5)
    for i, item in enumerate(stairs_len):
        a = random.random()
        if a < 0.5:
            step_height = -step_height
        stop_x = start_x + step_width
        stair_offset = random.choice([0.1, 0.15, -0.1, -0.15])
        end = (item + stair_offset) / terrain.horizontal_scale + start_x
        while (stop_x) < end:
            start_x += step_width
            stop_x += step_width
            height += step_height
            terrain.height_field_raw[start_x:stop_x, :] = height
        if i == len(stairs_len) - 1:
            terrain.height_field_raw[stop_x:, :] = height
        else:
            terrain.height_field_raw[stop_x:, :] = height
        start_x = stop_x + platform_len
    
    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height
    
def generate_obstacle(height_field_raw, virtual_height_field, virtual_reward_field, obs_pos, length, width, obs_height):
    rand_x = np.random.randint(5, 10)
    # rand_y = np.random.randint(10, 20)
    rand_y = 0
    
    obs_center = [0, 0]
    obs_center[0] = obs_pos[0] + rand_x
    obs_center[1] = obs_pos[1] + rand_y
    
    x = np.arange(length)
    y = np.arange(width)
    X, Y = np.meshgrid(x, y)
    
    center_x = length // 2
    center_y = width // 2

    obs_reward_map = np.sqrt((X-center_x)**2 + (Y-center_y)**2)
    obs_reward_map = 1/(obs_reward_map/obs_reward_map.max()+0.1)  # center:100 edge:1
    
    width_low = np.clip(int(obs_center[1]-width/2), a_min=0, a_max=virtual_reward_field.shape[1])
    width_high = np.clip(int(obs_center[1]+width/2), a_min=0, a_max=virtual_reward_field.shape[1])
    
    virtual_reward_field[int(obs_center[0]-length/2):int(obs_center[0]+length/2), 
                     width_low:width_high] = (obs_reward_map.T)[:, :width_high-width_low]
    
    real_width_low = np.clip(int(obs_center[1]-width/2), a_min=0, a_max=height_field_raw.shape[1]-1); real_width_high = np.clip(int(obs_center[1]+width/2), a_min=0, a_max=height_field_raw.shape[1]-1)
    real_length_low = np.clip(int(obs_center[0]-length/2), a_min=0, a_max=height_field_raw.shape[0]-1); real_length_high = np.clip(int(obs_center[0]+length/2), a_min=0, a_max=height_field_raw.shape[0]-1)
    height_field_raw[real_length_low:real_length_high, 
                     real_width_low:real_width_high] = obs_height
    # height_field_raw[int(obs_center[0]-length/2):int(obs_center[0]+length/2), 
    #                  int(obs_center[1]-width/2):int(obs_center[1]+width/2)] = obs_height
        
    # # virtual student
    # virtual_width_low = np.clip(real_width_low, a_min=5, a_max=height_field_raw.shape[1]-1); virtual_width_high = np.clip(real_width_high, a_min=0, a_max=virtual_height_field.shape[1]-6)
    # virtual_length_low = np.clip(real_length_low, a_min=3, a_max=virtual_height_field.shape[0]-1); virtual_length_high = np.clip(real_length_high, a_min=0, a_max=virtual_height_field.shape[0]-4)
    # virtual_height_field[virtual_length_low-3:virtual_length_high+3, 
    #                  virtual_width_low-5:virtual_width_high+5] = obs_height
    
    # virtual teacher
    virtual_width_low = np.clip(real_width_low, a_min=3, a_max=height_field_raw.shape[1]-1); virtual_width_high = np.clip(real_width_high, a_min=0, a_max=virtual_height_field.shape[1]-4)
    virtual_length_low = np.clip(real_length_low, a_min=2, a_max=virtual_height_field.shape[0]-1); virtual_length_high = np.clip(real_length_high, a_min=0, a_max=virtual_height_field.shape[0]-3)
    virtual_height_field[virtual_length_low-2:virtual_length_high+2, 
                     virtual_width_low-3:virtual_width_high+3] = obs_height
    
    return height_field_raw, virtual_height_field, virtual_reward_field

def virtual_flat_terrain(terrain,
                    platform_len=2.5, 
                    platform_height=0., 
                    num_stones=8,
                    obstacle_size=[1, 2, 3],
                    x_range=[1.5, 2.4],
                    y_range=[-0.4, 0.4],
                    pad_width=0.1,
                    pad_height=0.5,):
    goals = np.zeros((num_stones+2, 3+4))
    
    mid_y = terrain.length // 2  # length is actually y width

    dis_x_min = round(x_range[0] / terrain.horizontal_scale)
    dis_x_max = round(x_range[1] / terrain.horizontal_scale)
    dis_y_min = round(y_range[0] / terrain.horizontal_scale)
    dis_y_max = round(y_range[1] / terrain.horizontal_scale)
    
    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    terrain.height_field_raw[0:platform_len, :] = platform_height
    terrain.virtual_height_field[0:platform_len, :] = platform_height

    dis_x = platform_len
    goals[0, :2] = [platform_len - 1, mid_y]
    last_y = mid_y
    obs_type = random.random()
    # obs_type = 0.2
    if obs_type < 0.5:
        for i in range(num_stones):
            rand_x = np.random.randint(dis_x_min, dis_x_max)
            # import ipdb; ipdb.set_trace()
            rand_y = np.random.randint(dis_y_min, dis_y_max)
            if i == 1 or i == 3:
                rand_x *= 2.
                terrain.height_field_raw, terrain.virtual_height_field, terrain.virtual_reward_field = generate_obstacle(terrain.height_field_raw, 
                                                                terrain.virtual_height_field,
                                                                terrain.virtual_reward_field,
                                                                [dis_x+rand_x//2, (mid_y + rand_y + last_y)//2],
                                                                length=int(obstacle_size[0]/terrain.horizontal_scale), 
                                                                width=int(obstacle_size[1]/terrain.horizontal_scale), 
                                                                obs_height=obstacle_size[2]/terrain.vertical_scale)
            dis_x += rand_x
            # goals[i+1, :2] = [dis_x-rand_x//2, mid_y + rand_y]
            goals[i+1, :2] = [dis_x, mid_y + rand_y]
            last_y = mid_y + rand_y
    else:
        for i in range(num_stones):
            rand_x = np.random.randint(dis_x_min, dis_x_max)
            # import ipdb; ipdb.set_trace()
            rand_y = np.random.randint(dis_y_min, dis_y_max)
            if i == 2:
                rand_x *= 2.
                terrain.height_field_raw, terrain.virtual_height_field, terrain.virtual_reward_field = generate_obstacle(terrain.height_field_raw, 
                                                                terrain.virtual_height_field,
                                                                terrain.virtual_reward_field,
                                                                [dis_x+rand_x//2, (mid_y + rand_y + last_y)//2],
                                                                length=int(obstacle_size[0]/terrain.horizontal_scale), 
                                                                width=int(obstacle_size[1]/terrain.horizontal_scale), 
                                                                obs_height=obstacle_size[2]/terrain.vertical_scale)
            dis_x += rand_x
            # goals[i+1, :2] = [dis_x-rand_x//2, mid_y + rand_y]
            goals[i+1, :2] = [dis_x, mid_y + rand_y]
            last_y = mid_y + rand_y
    
    final_dis_x = dis_x + (np.random.randint(dis_x_min, dis_x_max)//2)

    if final_dis_x > terrain.width:
        final_dis_x = terrain.width - 0.5 // terrain.horizontal_scale
    goals[-1, :2] = [final_dis_x, mid_y]
    
    terrain.goals = goals * terrain.horizontal_scale
    
    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height; terrain.virtual_height_field[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height; terrain.virtual_height_field[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height; terrain.virtual_height_field[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height; terrain.virtual_height_field[-pad_width:, :] = pad_height

def generate_half_obstacle(height_field_raw, virtual_height_field, virtual_reward_field, flat_height, gap_size, obs_pos, length, width, obs_height):
    rand_x = np.random.randint(5, 10)
    # rand_y = np.random.randint(10, 20)
    rand_y = 0
    
    obs_center = [0, 0]
    obs_center[0] = obs_pos[0] + rand_x
    obs_center[1] = obs_pos[1] + rand_y
    
    x = np.arange(length)
    y = np.arange(width)
    X, Y = np.meshgrid(x, y)
    
    center_x = length // 2
    center_y = width // 2

    obs_reward_map = np.sqrt((X-center_x)**2 + (Y-center_y)**2)
    obs_reward_map = 1/(obs_reward_map/obs_reward_map.max()+0.1)  # center:100 edge:1
    
    width_low = np.clip(int(obs_center[1]-width/2), a_min=0, a_max=virtual_reward_field.shape[1])
    width_high = np.clip(int(obs_center[1]+width/2), a_min=0, a_max=virtual_reward_field.shape[1])
    
    virtual_reward_field[int(obs_center[0]-length/2):int(obs_center[0]+length/2), 
                     width_low:width_high] = (obs_reward_map.T)[:, :width_high-width_low]
    
    # virtual_width = 3 + 2; virtual_length = 2 + 1  # student
    virtual_width = 3; virtual_length = 2  # teacher
    left_or_right = random.random(); flat_type = random.random()
    flat_height = flat_height if flat_type < 0.5 else -obs_height // 3
    
    off_set = random.randint(-4, 4); off_center_1 = int(obs_center[1]+off_set)
    if left_or_right < 0.5:
        real_off_center_1 = np.clip(off_center_1, a_min=0, a_max=height_field_raw.shape[1]-1)
        virtual_off_center_1 = np.clip(off_center_1-virtual_width, a_min=0, a_max=height_field_raw.shape[1]-1)
        height_field_raw[int(obs_center[0]-length/2):int(obs_center[0]+length/2), real_off_center_1:] = obs_height
        virtual_height_field[int(obs_center[0]-length/2)-virtual_length:int(obs_center[0]+length/2)+virtual_length, virtual_off_center_1:] = obs_height
        if flat_height < 0:
            height_field_raw[int(obs_center[0]-gap_size/2):int(obs_center[0]+gap_size/2), :real_off_center_1] = flat_height
            virtual_height_field[int(obs_center[0]-gap_size/2):int(obs_center[0]+gap_size/2), :virtual_off_center_1] = flat_height
        else:
            height_field_raw[int(obs_center[0]-length/2):int(obs_center[0]+length/2), :real_off_center_1] = flat_height
            virtual_height_field[int(obs_center[0]-length/2):int(obs_center[0]+length/2), :virtual_off_center_1] = flat_height
    else:
        real_off_center_1 = np.clip(off_center_1, a_min=0, a_max=height_field_raw.shape[1]-1)
        virtual_off_center_1 = np.clip(off_center_1+virtual_width, a_min=0, a_max=height_field_raw.shape[1]-1)
        height_field_raw[int(obs_center[0]-length/2):int(obs_center[0]+length/2), :real_off_center_1] = obs_height
        virtual_height_field[int(obs_center[0]-length/2)-virtual_length:int(obs_center[0]+length/2)+virtual_length, :virtual_off_center_1] = obs_height
        if flat_height < 0:
            height_field_raw[int(obs_center[0]-gap_size/2):int(obs_center[0]+gap_size/2), real_off_center_1:] = flat_height
            virtual_height_field[int(obs_center[0]-gap_size/2):int(obs_center[0]+gap_size/2), virtual_off_center_1:] = flat_height
        else:
            height_field_raw[int(obs_center[0]-length/2):int(obs_center[0]+length/2), real_off_center_1:] = flat_height
            virtual_height_field[int(obs_center[0]-length/2):int(obs_center[0]+length/2), virtual_off_center_1:] = flat_height        
    
    return height_field_raw, virtual_height_field, virtual_reward_field

def half_virtual_terrain(terrain,
                    platform_len=2.5, 
                    platform_height=0., 
                    num_stones=8,
                    obstacle_size=[1, 2, 3],
                    x_range=[1.5, 2.4],
                    y_range=[-0.4, 0.4],
                    pad_width=0.1,
                    pad_height=0.5,
                    hurdle_height_range=[0.5, 1.5],
                    gap_size=0.1):
    goals = np.zeros((num_stones+2, 3+4))
    hurdle_height_max = round(hurdle_height_range[1] / terrain.vertical_scale)
    hurdle_height_min = round(hurdle_height_range[0] / terrain.vertical_scale)
    hurdle_height = np.random.randint(hurdle_height_min, hurdle_height_max)
    
    mid_y = terrain.length // 2  # length is actually y width

    dis_x_min = round(x_range[0] / terrain.horizontal_scale)
    dis_x_max = round(x_range[1] / terrain.horizontal_scale)
    dis_y_min = round(y_range[0] / terrain.horizontal_scale)
    dis_y_max = round(y_range[1] / terrain.horizontal_scale)
    
    platform_len = round(platform_len / terrain.horizontal_scale)
    platform_height = round(platform_height / terrain.vertical_scale)
    terrain.height_field_raw[0:platform_len, :] = platform_height
    terrain.virtual_height_field[0:platform_len, :] = platform_height

    dis_x = platform_len
    goals[0, :2] = [platform_len - 1, mid_y]
    last_y = mid_y 
    obs_type = random.random()  
    # obs_type = 0.2
    # gap_size = 1.0
    if obs_type < 0.5:
        for i in range(num_stones):
            rand_x = np.random.randint(dis_x_min, dis_x_max)
            # import ipdb; ipdb.set_trace()
            rand_y = np.random.randint(dis_y_min, dis_y_max)
            if i == 1 or i == 3:
                rand_x *= 2.
                terrain.height_field_raw, terrain.virtual_height_field, terrain.virtual_reward_field = generate_half_obstacle(terrain.height_field_raw, 
                                                                terrain.virtual_height_field,
                                                                terrain.virtual_reward_field,
                                                                hurdle_height,
                                                                gap_size/terrain.horizontal_scale,
                                                                [dis_x+rand_x//2, (mid_y + rand_y + last_y)//2],
                                                                length=int(obstacle_size[0]/terrain.horizontal_scale), 
                                                                width=int(obstacle_size[1]/terrain.horizontal_scale), 
                                                                obs_height=obstacle_size[2]/terrain.vertical_scale)
            dis_x += rand_x
            goals[i+1, :2] = [dis_x, mid_y + rand_y]
            last_y = mid_y + rand_y
    else:
        for i in range(num_stones):
            rand_x = np.random.randint(dis_x_min, dis_x_max)
            rand_y = np.random.randint(dis_y_min, dis_y_max)
            if i == 2:
                rand_x *= 2.
                terrain.height_field_raw, terrain.virtual_height_field, terrain.virtual_reward_field = generate_half_obstacle(terrain.height_field_raw, 
                                                                terrain.virtual_height_field,
                                                                terrain.virtual_reward_field,
                                                                hurdle_height,
                                                                gap_size/terrain.horizontal_scale,
                                                                [dis_x+rand_x//2, (mid_y + rand_y + last_y)//2],
                                                                length=int(obstacle_size[0]/terrain.horizontal_scale), 
                                                                width=int(obstacle_size[1]/terrain.horizontal_scale), 
                                                                obs_height=obstacle_size[2]/terrain.vertical_scale)
            dis_x += rand_x
            goals[i+1, :2] = [dis_x, mid_y + rand_y]
            last_y = mid_y + rand_y
    
    final_dis_x = dis_x + (np.random.randint(dis_x_min, dis_x_max)//2)

    if final_dis_x > terrain.width:
        final_dis_x = terrain.width - 0.5 // terrain.horizontal_scale
    goals[-1, :2] = [final_dis_x, mid_y]
    
    terrain.goals = goals * terrain.horizontal_scale
    
    # pad edges
    pad_width = int(pad_width // terrain.horizontal_scale)
    pad_height = int(pad_height // terrain.vertical_scale)
    terrain.height_field_raw[:, :pad_width] = pad_height; terrain.virtual_height_field[:, :pad_width] = pad_height
    terrain.height_field_raw[:, -pad_width:] = pad_height; terrain.virtual_height_field[:, -pad_width:] = pad_height
    terrain.height_field_raw[:pad_width, :] = pad_height; terrain.virtual_height_field[:pad_width, :] = pad_height
    terrain.height_field_raw[-pad_width:, :] = pad_height; terrain.virtual_height_field[-pad_width:, :] = pad_height
